package com.dianwang.utils;

import com.huaban.analysis.jieba.JiebaSegmenter;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.*;

/**
 * 文本预处理类，包含停用词过滤和同义词替换功能。
 */
public class TextPreprocessor {
    private static final Set<String> STOP_WORDS = new HashSet<>(Arrays.asList(
            "的", "了", "在", "是", "和", "有", "与", "为", "于", "或", "但", "因此", "为此"
    ));

    public static final Map<String, String> SYNONYM_MAP = new HashMap<>();

    static {
        loadSynonyms("synonyms.txt");
    }

    /**
     * 从指定的同义词文件中加载同义词映射。
     *
     * @param synonymFile 同义词文件名，位于资源目录下
     */
    private static void loadSynonyms(String synonymFile) {
        try (InputStream is = TextPreprocessor.class.getClassLoader().getResourceAsStream(synonymFile);
             BufferedReader reader = new BufferedReader(new InputStreamReader(is, "UTF-8"))) {

            String line;
            while ((line = reader.readLine()) != null) {
                // 每行一个同义词组，逗号分隔
                String[] synonyms = line.split(",");
                if (synonyms.length < 2) continue; // 至少需要两个同义词

                // 将所有同义词映射到第一个词
                String primary = synonyms[0].trim();
                for (int i = 1; i < synonyms.length; i++) {
                    SYNONYM_MAP.put(synonyms[i].trim(), primary);
                }
            }
            System.out.println("同义词加载完成，共加载 " + SYNONYM_MAP.size() + " 个同义词。");
        } catch (Exception e) {
            System.out.println("加载同义词文件失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 预处理输入文本，包括：
     * 1. 替换同义词
     * 2. 去除停用词
     * 3. 去除标点符号
     * 4. 分词
     *
     * @param input 输入文本
     * @return 预处理后的文本
     */
    private static final String USER_DICT_PATH = "userdict.txt";

    public static String preprocess(String input) {
        if (input == null) return "";

        // 替换 KV => kV
        input = input.replaceAll("KV", "kV").replaceAll("kv", "kV");

        // 去除标点符号
        input = input.replaceAll("[，。！？、；：]", " ");

        // 分词（使用结巴分词）
        JiebaSegmenter segmenter = new JiebaSegmenter();
        // 加载用户词典
        try {
            java.net.URL dictUrl = TextPreprocessor.class.getClassLoader().getResource(USER_DICT_PATH);
            if (dictUrl == null) {
                throw new RuntimeException("无法找到用户词典文件: " + USER_DICT_PATH);
            }
            com.huaban.analysis.jieba.WordDictionary.getInstance().loadUserDict(java.nio.file.Paths.get(dictUrl.toURI()));
        } catch (Exception e) {
            System.out.println("警告：加载用户词典失败: " + e.getMessage());
        }
        List<String> tokens = new ArrayList<>();
        for (com.huaban.analysis.jieba.SegToken token : segmenter.process(input, JiebaSegmenter.SegMode.INDEX)) {
            tokens.add(token.word);
        }

        // 去除停用词和同义词替换
        List<String> processedTokens = new ArrayList<>();
        for (String token : tokens) {
            if (!STOP_WORDS.contains(token)) {
                // 替换同义词
                String synonym = SYNONYM_MAP.getOrDefault(token, token);
                processedTokens.add(synonym);
            }
        }

        return String.join(" ", processedTokens);
    }

    public static void main(String[] args) {
        String testInput = "巡视发现：#2主变有载调压呼吸器底部油杯内油位较低，接近最低下限主变压器\n" + "加热器\n" + "加热器损坏\n" + "油封油过少。";
        String processed = preprocess(testInput);
        System.out.println(processed);
    }
}
